Skip to content

CUDA support in blocks - #1599

Merged
NicolasHug merged 17 commits into
meta-pytorch:mainfrom
NicolasHug:blocks_cuda_mine
Aug 5, 2026
Merged

CUDA support in blocks#1599
NicolasHug merged 17 commits into
meta-pytorch:mainfrom
NicolasHug:blocks_cuda_mine

Conversation

@NicolasHug

@NicolasHug NicolasHug commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

This enables our NVDEC CUDA interface on the "Blocks" APIs. Lots of things are still not supported but the basic functionality works and is on par to CPU.

The original implementation was vibe-coded in #1592, and I re-implemented everything here from scratch because these aren't trivial changes and I really needed to understand them.

A few key points:

  • The CUDA frame is copied by the PacketDecoder in a new make_frame_standalone() method before it is returned (and potentially later color-converted). That is needed to satisfy an assumption of the NVDEC decoding loop: we always unmap just before we map a new frame, and for that to be correct, the frame that we unmap must not be needed anymore. In SingleStreamDecoder mode, we rely on the fact that this frame (that we unmap) is either discarded (we don't care about it) or color-converted (and thus copied), which makes it unmap-able. Here, we cannot assume that anymore, so we enforce a manual copy.
  • The color-conversion entry point (convert_av_frame_to_frame_output()) was previously relying on the interface's state - a state that came from video decoding parts, which aren't present in its color-convert-only mode. This PR makes convert_av_frame_to_frame_output() rely on info that doesn't come from that state. Particularly, part of that state is now attached to the AVFrame via the StandAloneFrameAttachedData struct.
  • Now that the interface entry points can be accessed from different threads, we need a more robust CUDA context management: see the new CudaContextGuard and associated comment.

Everything downstream of an AVFrame's owner took `UniqueAVFrame&` or
`const UniqueAVFrame&`. That's a constraint on the caller's storage rather
than a statement about what the function does, and it has two costs.

A non-const `UniqueAVFrame&` lets a callee take ownership of the caller's
frame. Both CUDA interfaces did: BetaCudaDeviceInterface moved out of it,
CudaDeviceInterface reassigned it. That is invisible under
SingleStreamDecoder, whose frame is a loop local that dies right after
conversion, but the building-block ops own their frame in a handle that
outlives the call, so the frame gets freed twice.

`const UniqueAVFrame&` doesn't allow the steal, but it still forces anyone
holding a plain AVFrame* -- which is what the ops' tensor handle really is
-- to manufacture a unique_ptr just to make the call, and manufacturing a
second owner for an already-owned object is its own bug factory.

So: functions that only look at a frame now take `const AVFrame&`, and the
one that writes to it takes `AVFrame&`. Ownership stays with whoever
actually owns the frame. Producers (receive_frame) keep `UniqueAVFrame&`
because they really do hand back ownership.

With that, the ops layer needs no ownership sleight-of-hand:
wrap_pointer_to_tensor() gains a deleter parameter, so the one generic
handle covers Demuxer/PacketDecoder/ColorConverter and the FFmpeg types
too, and both bespoke wrap_*_pointer_to_tensor() functions go away.
Demuxer::next_packet() returns UniqueAVPacket rather than a raw pointer
plus a comment telling the caller to free it.

The encoder is left alone: it owns and mutates its frames, and it uses a
null frame as the flush signal, so a reference is the wrong shape there.
Two things only show up when building against the older FFmpeg headers.

get_num_channels() was patching av_frame.channel_layout when FFmpeg 4 left
it unset, so it was a mutator wearing an observer's name, and the layout
fix-up was a side effect that swresample setup silently relied on. Pull the
fix-up into get_channel_layout() and call that from the two places that
actually need a layout; get_num_channels() just counts.

swr_alloc_set_opts2() only became const-correct in FFmpeg 6 (libswresample
4.12). Cast for the older headers, which don't modify the layout either.
@pytorch-bot

pytorch-bot Bot commented Aug 4, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/meta-pytorch/torchcodec/1599

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 4, 2026
@NicolasHug
NicolasHug marked this pull request as ready for review August 5, 2026 16:09
@NicolasHug
NicolasHug merged commit 64763fc into meta-pytorch:main Aug 5, 2026
13 checks passed
@NicolasHug
NicolasHug deleted the blocks_cuda_mine branch August 5, 2026 16:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant